# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: daniel

import os
import matplotlib.pyplot as plt
import pandas as pd
import pathlib
import jinja2
import collections
from absl import app, flags

import coin.proto.coin_order_enums_pb2 as coin_enum


FLAGS = flags.FLAGS


def _sumzip(*items):
  return [sum(values) for values in zip(*items)]


class FailOrdersAnalyzer(object):
  def __init__(self):
    self._fail_orders_summary = []
    self._error_stats_summary = []

  def filter_fail_order_records(self, records):
    ret = []
    for record in records:
      if record['type'] != 'OG_LOG' or record['og_log']['type'] != 'ORDER_EVENT':
        continue
      if record['og_log']['event']['type'] not in ['ORDER_REJECTED', 'CANCEL_REJECTED']:
        continue
      raw_msg_id = int(record['og_log']['raw_msg_id'])
      source = coin_enum.OrderSubTopicId.Name(raw_msg_id >> (64 - 8))
      if source in ['ORDER_EXECUTION', 'SYSTEM']:
        continue
      ret.append(record)
    return ret

  def retrieve_fail_orders_record(self,
                                  records,
                                  action,
                                  start_dt,
                                  end_dt):
    tmp = []
    for record in records:
      if record['og_log']['event']['type'] != action:
        continue
      rec = {
          'timestamp': int(record['og_log']['event']['event_time']),
          'num_of_hit': 1,
      }
      tmp.append(rec)
    num = len(tmp)
    ret = []
    x_beg = int(pd.to_datetime(start_dt.strftime('%Y%m%dT%H%M%S')).timestamp()) * 1000000000
    x_end = int(pd.to_datetime(end_dt.strftime('%Y%m%dT%H%M%S')).timestamp()) * 1000000000
    x = x_beg
    y = 0
    while x_beg < x_end:
      for each in tmp:
        if each['timestamp'] < x_beg:
          continue
        if each['timestamp'] >= x_beg + 300000000000:
          continue
        y = y + 1
      info = {'timestamp': x, 'num_of_hit': y}
      ret.append(info)
      x_beg = x_beg + 300000000000
      x = x_beg
      y = 0

    return num, sorted(ret, key=lambda x: x['timestamp'])

  def generate_fail_orders_summary_report(self,
                                          trade_date,
                                          output_dir,
                                          local_run,
                                          slurm_run):
    if slurm_run:
      # another process flow for slurm run
      return
    if len(self._fail_orders_summary) == 0:
      return

    plt.rcParams.update({'font.size': 7})
    fig = plt.figure(figsize=(30, 16), dpi=200)
    # draw bar chart
    records = sorted(self._fail_orders_summary, key=lambda x: x['total_num'], reverse=True)
    color = ['steelblue', 'orange']
    label = []
    order_rejected_color = []
    order_rejected_total = []
    cancel_rejected_color = []
    cancel_rejected_total = []
    for record in records:
      order_rejected_num = int(record['order_rejected_num'])
      cancel_rejected_num = int(record['cancel_rejected_num'])
      total = order_rejected_num + cancel_rejected_num
      if total == 0:
        continue
      label.append(record['strategy_name'])
      order_rejected_total.append(order_rejected_num)
      order_rejected_color.append(color[0])
      cancel_rejected_total.append(cancel_rejected_num)
      cancel_rejected_color.append(color[1])

    plt.subplot(211)
    plt.bar(label, order_rejected_total, width=0.35, label='Order Rejected',
            color=order_rejected_color)
    plt.bar(label, cancel_rejected_total, width=0.35, label='Cancel Rejected',
            color=cancel_rejected_color, bottom=_sumzip(order_rejected_total))
    plt.ylabel('total # of failed orders')
    plt.xticks(range(len(label)), label, rotation=270)
    plt.title('failed orders summary')
    # draw bar chart
    records = []
    for each in self._error_stats_summary:
      api_limit_error_num = 0
      error_stats = each['error_stats']
      if 'submit' in error_stats and \
         'error_code' in error_stats['submit'] and \
         'API_LIMIT_ERROR' in error_stats['submit']['error_code']:
        api_limit_error_num += int(error_stats['submit']['error_code']['API_LIMIT_ERROR'])
      if 'cancel' in error_stats and \
         'error_code' in error_stats['cancel'] and \
         'API_LIMIT_ERROR' in error_stats['cancel']['error_code']:
        api_limit_error_num += int(error_stats['cancel']['error_code']['API_LIMIT_ERROR'])
      if api_limit_error_num == 0:
        continue
      info = {
          "strategy_name": each["strategy_name"],
          "api_limit_error_num": api_limit_error_num,
      }
      records.append(info)
    records = sorted(records, key=lambda x: x['api_limit_error_num'], reverse=True)
    label = []
    total = []
    if len(records) == 0:
      return
    for record in records:
      label.append(record['strategy_name'])
      total.append(record['api_limit_error_num'])
    plt.subplot(212)
    plt.bar(label, total, width=0.35, label='API LIMIT ERROR')
    plt.ylabel('total # of api limit error')
    plt.xticks(range(len(label)), label, rotation=270)
    plt.title('api limit error summary')

    plt.grid(True)
    plt.legend()
    fig.tight_layout()
    if local_run is True:
      plt.show()
    else:
      output_path = f"{output_dir}/{trade_date}/fail_orders"
      if not os.path.exists(output_path):
        os.makedirs(output_path)
      output_file = f"{trade_date}_fail_orders_summary_report.svg"
      print(f"write to {output_path}/{output_file}")
      fig.savefig(f"{output_path}/{output_file}", dpi=fig.dpi)
    plt.close(fig)

  def collect_error_stats(self, records, strategy_name):
    def item_construct():
      def item_construct2():
        return collections.defaultdict(lambda: 0)
      return collections.defaultdict(item_construct2)

    error_stats = collections.defaultdict(item_construct)
    for record in records:
      event = record['og_log']['event']

      error_code = event.get('error_code')
      etype = 'submit' if event['type'] == 'ORDER_REJECTED' else 'cancel'

      if error_code is None:
        print("No error code in event", event)
      else:
        error_code = coin_enum.OrderErrorCode.Name(error_code)
        error_stats[etype]['error_code'][error_code] += 1

      exchange_error_code = event.get('exchange_error_code')
      if exchange_error_code is None:
        print("No exchange error code in event", event)
      else:
        error_stats[etype]['exchange_error_code'][exchange_error_code] += 1

      exchange_error_msg = event.get('exchange_error_msg')
      if exchange_error_code is None:
        print("No exchange error msg in event", event)
      else:
        error_stats[etype]['exchange_error_msg'][exchange_error_msg] += 1

    info = {
        "strategy_name": strategy_name,
        "error_stats": error_stats,
    }
    self._error_stats_summary.append(info)
    # prepare srun scripts
    template_loader = jinja2.FileSystemLoader(searchpath=os.path.dirname(__file__))
    template_env = jinja2.Environment(loader=template_loader)
    template = template_env.get_template('error_stats.html.jinja')
    html = template.render(strategy_name=strategy_name, error_stats=error_stats)
    return html

  def generate_fail_orders_report(self,
                                  records,
                                  strategy_name,
                                  machine,
                                  start_dt,
                                  end_dt,
                                  output_dir,
                                  local_run):
    fail_order_records = self.filter_fail_order_records(records)
    # ignore strategy with zero failed order
    if len(fail_order_records) == 0:
      return

    order_rejected_num, order_rejected_records = \
        self.retrieve_fail_orders_record(fail_order_records, 'ORDER_REJECTED', start_dt, end_dt)
    cancel_rejected_num, cancel_rejected_records = \
        self.retrieve_fail_orders_record(fail_order_records, 'CANCEL_REJECTED', start_dt, end_dt)
    html = self.collect_error_stats(fail_order_records, strategy_name)

    fail_order_total_num = order_rejected_num + cancel_rejected_num
    name = strategy_name
    info = {
        'strategy_name': name,
        'total_num': fail_order_total_num,
        'order_rejected_num': order_rejected_num,
        'cancel_rejected_num': cancel_rejected_num,
    }
    self._fail_orders_summary.append(info)

    label = ['Order_Rejected', 'Cancel_Rejected']
    color = ['steelblue', 'orange']
    total = [
        order_rejected_num,
        cancel_rejected_num,
    ]
    explode = [0, 0]
    plt.rcParams.update({'font.size': 7})
    fig = plt.figure(figsize=(10, 6), dpi=200)
    # draw pie chart
    plt.subplot(221)
    plt.pie(total, colors=color, explode=explode, labels=label, shadow=True, autopct='%1.1f%%')

    # draw bar chart
    plt.subplot(222)
    plt.bar(label, total, width=0.5, label='', color=color)
    plt.ylabel('# of hits')

    # draw linear chart
    plt.subplot(212)
    if order_rejected_num > 0:
      df = pd.DataFrame(order_rejected_records)
      df.timestamp = pd.to_datetime(df.timestamp, unit='ns')
      plt.plot(df.timestamp,
               df.num_of_hit,
               marker='.',
               markersize=0.1,
               label='Order_Rejected',
               color=color[0])

    if cancel_rejected_num > 0:
      df = pd.DataFrame(cancel_rejected_records)
      df.timestamp = pd.to_datetime(df.timestamp, unit='ns')
      plt.plot(df.timestamp,
               df.num_of_hit,
               marker='.',
               markersize=0.1,
               label='Cancel_Rejected',
               color=color[1])

    plt.title('# of fail order in 5 mins')
    plt.grid(True)
    plt.legend()
    fig.tight_layout(pad=1.0, w_pad=1.0, h_pad=2.0)
    if local_run is True:
      plt.show()
    else:
      trade_date = start_dt.strftime('%Y%m%d')
      start_dt_str = start_dt.strftime('%Y%m%dT%H%M%S')
      end_dt_str = end_dt.strftime('%Y%m%dT%H%M%S')
      output_path = f"{output_dir}/{trade_date}/fail_orders/{strategy_name}/{machine}"
      if not os.path.exists(output_path):
        os.makedirs(output_path)
      output_file = f"{start_dt_str}-{end_dt_str}.svg"
      print(f"write to {output_path}/{output_file}")
      fig.savefig(f"{output_path}/{output_file}", dpi=fig.dpi)
    plt.close(fig)

    htmlpath = pathlib.Path(f"{output_path}/error_stats.html")
    htmlpath.write_text(html)


def main():
  # TODO(daniel):
  print("@daniel: to impliment")


if __name__ == '__main__':
  flags.DEFINE_string('root_dir',
                      '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log',
                      'root dir specify')
  flags.DEFINE_string('output_dir',
                      '/remote/iosg/strat-1/buckets/analysis.derived.coin/live/plot/og_log_plot',
                      'output_dir.')
  flags.DEFINE_string('tmp_dir',
                      '/tmp',
                      'tmp dir for slurm run')
  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')
  flags.DEFINE_string('machine', None, 'Instance machine name.')
  flags.DEFINE_string('strategy_name', None, 'Strategy name.')
  flags.DEFINE_string('start_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('end_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('symbol', None, '')
  flags.DEFINE_string('log_type', None, '')
  flags.DEFINE_bool('local_run', False, "")
  flags.DEFINE_bool('slurm_run', False, "")
  flags.DEFINE_bool('generate_summary_report', False, "generate summary report only")
  app.run(main)
